
import json
import re
import os
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
from pathlib import Path
import traceback
from collections import defaultdict
from sklearn.metrics import precision_recall_fscore_support, classification_report, confusion_matrix, roc_curve, auc
from transformers import AutoTokenizer, AutoModelForTokenClassification, TrainingArguments, Trainer
from datasets import Dataset
import matplotlib.pyplot as plt
import seaborn as sns

# 确保模型目录存在
model_dir = Path("/root/npl/data/biobert-v1.1")
model_dir.mkdir(parents=True, exist_ok=True)

print(f"检测到本地模型文件完整，目录: {model_dir.resolve()}")

# 定义简单的生成器模型
class Generator(nn.Module):
    def __init__(self, vocab_size, hidden_dim, max_length):
        super(Generator, self).__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.lstm = nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, vocab_size)
        self.max_length = max_length

    def forward(self, x):
        embedded = self.embedding(x)
        lstm_out, _ = self.lstm(embedded)
        logits = self.fc(lstm_out)
        return logits

# 定义简单的判别器模型
class Discriminator(nn.Module):
    def __init__(self, vocab_size, hidden_dim):
        super(Discriminator, self).__init__()
        self.embedding = nn.Embedding(vocab_size, hidden_dim)
        self.lstm = nn.LSTM(hidden_dim, hidden_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, 1)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        embedded = self.embedding(x)
        lstm_out, _ = self.lstm(embedded)
        out = self.fc(lstm_out[:, -1, :])
        return self.sigmoid(out)

def clean_dataset(file_content: str, dataset_name: str):
    """
    精简版数据集清洗函数
    """
    documents = []
    entity_labels = set()
    
    for i, line in enumerate(file_content.strip().split('\n')):
        line = line.strip()
        if not line:
            continue
            
        if '\t' not in line:
            continue
            
        parts = line.split('\t', 1)
        doc_id = parts[0].strip()
        raw_text = parts[1].strip()
        
        if not doc_id or not raw_text:
            continue
            
        entities = []
        for match in re.finditer(r'<category="([^"]+)">([^<]+)</category>', raw_text):
            label = match.group(1)
            entity_text = match.group(2)
            start = match.start()
            end = match.end()
            entities.append({
                "label": label,
                "text": entity_text,
                "start": start,
                "end": end
            })
            entity_labels.add(label)
        
        documents.append({
            "id": doc_id,
            "text": raw_text,
            "entities": entities
        })
        
    print(f"{dataset_name}清洗完成: 有效样本={len(documents)}, 标签类型={entity_labels}")
    return documents, entity_labels

def align_labels_with_tokens(text, entities, tokenized, label_map):
    """
    精确的标签对齐函数
    """
    offset_mapping = tokenized["offset_mapping"]
    labels = [label_map["O"]] * len(offset_mapping)
    
    char_to_token = {}
    for token_idx, (start, end) in enumerate(offset_mapping):
        if start is None or end is None:
            continue
        for char_idx in range(int(start), int(end)):
            char_to_token[char_idx] = token_idx
    
    for entity in entities:
        entity_label = entity["label"]
        start_char = entity["start"]
        end_char = entity["end"]
        
        entity_tokens = set()
        for char_idx in range(start_char, end_char):
            if char_idx in char_to_token:
                entity_tokens.add(char_to_token[char_idx])
        
        if not entity_tokens:
            continue
        
        sorted_tokens = sorted(entity_tokens)
        for i, token_idx in enumerate(sorted_tokens):
            if i == 0:
                labels[token_idx] = label_map[f"B-{entity_label}"]
            else:
                labels[token_idx] = label_map[f"I-{entity_label}"]
    
    return labels

def tokenize_and_align_labels(examples, tokenizer, label_map):
    """
    令牌化文本并精确对齐标签
    """
    tokenized_inputs = tokenizer(
        examples["text"], 
        padding="max_length", 
        truncation=True, 
        max_length=128,
        return_offsets_mapping=True
    )
    
    labels = []
    for i in range(len(examples["text"])):
        text = examples["text"][i]
        entities = examples["entities"][i]
        tokenized = {
            "input_ids": tokenized_inputs["input_ids"][i],
            "attention_mask": tokenized_inputs["attention_mask"][i],
            "offset_mapping": tokenized_inputs["offset_mapping"][i]
        }
        label_seq = align_labels_with_tokens(text, entities, tokenized, label_map)
        labels.append(label_seq)
    
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

def compute_metrics(p):
    """
    计算评价指标：精确率、召回率、F1分数
    """
    predictions, labels = p.predictions, p.label_ids
    predictions = np.argmax(predictions, axis=2)
    
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for l in label if l != -100]
        for label in labels
    ]
    
    flat_true_labels = [label for sublist in true_labels for label in sublist]
    flat_predictions = [label for sublist in true_predictions for label in sublist]
    
    overall_precision, overall_recall, overall_f1, _ = precision_recall_fscore_support(
        flat_true_labels, flat_predictions, average='macro', zero_division=0
    )
    
    accuracy = np.mean([p == l for p, l in zip(flat_predictions, flat_true_labels)])
    
    class_report = classification_report(
        flat_true_labels, 
        flat_predictions, 
        output_dict=True,
        zero_division=0
    )
    
    results = {
        "overall_precision": overall_precision,
        "overall_recall": overall_recall,
        "overall_f1": overall_f1,
        "overall_accuracy": accuracy
    }
    
    for label, metrics in class_report.items():
        if label.isdigit():
            label_name = id2label[int(label)]
            results[f"{label_name}_precision"] = metrics["precision"]
            results[f"{label_name}_recall"] = metrics["recall"]
            results[f"{label_name}_f1"] = metrics["f1-score"]
    
    return results

def compute_confusion_matrix_and_roc(p, label_list, id2label, dataset_name):
    """
    计算混淆矩阵和ROC曲线
    """
    predictions = np.argmax(p.predictions, axis=2)
    labels = p.label_ids
    
    true_predictions = [
        [label_list[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    true_labels = [
        [label_list[l] for l in label if l != -100]
        for label in labels
    ]
    
    flat_true_labels = [label for sublist in true_labels for label in sublist]
    flat_predictions = [label for sublist in true_predictions for label in sublist]
    
    # 计算混淆矩阵
    cm = confusion_matrix(flat_true_labels, flat_predictions, labels=label_list)
    
    # 保存混淆矩阵图
    plt.figure(figsize=(10, 8))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=label_list, yticklabels=label_list)
    plt.xlabel('预测标签')
    plt.ylabel('真实标签')
    plt.title(f'{dataset_name} 混淆矩阵')
    plt.savefig(f'./results/{dataset_name}_confusion_matrix.png')
    plt.close()
    
    # 计算ROC曲线（针对每个类别）
    from sklearn.preprocessing import label_binarize
    from itertools import cycle
    
    # 将标签二值化
    y_true_bin = label_binarize(flat_true_labels, classes=label_list)
    y_pred_bin = label_binarize(flat_predictions, classes=label_list)
    
    # 计算每个类别的ROC曲线和AUC
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    for i in range(len(label_list)):
        fpr[i], tpr[i], _ = roc_curve(y_true_bin[:, i], y_pred_bin[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])
    
    # 绘制ROC曲线
    plt.figure(figsize=(10, 8))
    colors = cycle(['aqua', 'darkorange', 'cornflowerblue', 'green', 'red', 'purple'])
    for i, color in zip(range(len(label_list)), colors):
        plt.plot(fpr[i], tpr[i], color=color, lw=2,
                 label=f'ROC曲线 (类别 {label_list[i]}, AUC = {roc_auc[i]:.2f})')
    
    plt.plot([0, 1], [0, 1], 'k--', lw=2)
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('假阳性率')
    plt.ylabel('真阳性率')
    plt.title(f'{dataset_name} 多类别ROC曲线')
    plt.legend(loc="lower right")
    plt.savefig(f'./results/{dataset_name}_roc_curve.png')
    plt.close()
    
    return cm, roc_auc

def generate_synthetic_data(generator, tokenizer, num_samples, max_length, vocab_size, device):
    """
    使用生成器生成合成数据
    """
    generator.eval()
    synthetic_texts = []
    synthetic_entities = []
    
    for _ in range(num_samples):
        noise = torch.randint(0, vocab_size, (1, max_length)).to(device)
        logits = generator(noise)
        probs = torch.softmax(logits, dim=-1)
        generated = torch.multinomial(probs.squeeze(0), num_samples=1).squeeze(-1)
        
        text = tokenizer.decode(generated, skip_special_tokens=True)
        
        entities = []
        for match in re.finditer(r'(\w+)', text):
            if np.random.rand() > 0.7:
                entities.append({
                    "label": np.random.choice(list(all_labels)),
                    "text": match.group(1),
                    "start": match.start(),
                    "end": match.end()
                })
        
        synthetic_texts.append(text)
        synthetic_entities.append(entities)
    
    return synthetic_texts, synthetic_entities

def train_gan(generator, discriminator, train_texts, tokenizer, num_epochs, device):
    """
    训练GAN模型
    """
    g_optimizer = optim.Adam(generator.parameters(), lr=1e-3)
    d_optimizer = optim.Adam(discriminator.parameters(), lr=1e-3)
    criterion = nn.BCELoss()
    
    for epoch in range(num_epochs):
        for text in train_texts:
            input_ids = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=128)["input_ids"].to(device)
            
            d_optimizer.zero_grad()
            real_labels = torch.ones(1, 1).to(device)
            real_output = discriminator(input_ids)
            d_real_loss = criterion(real_output, real_labels)
            
            noise = torch.randint(0, tokenizer.vocab_size, (1, 128)).to(device)
            gen_logits = generator(noise)
            gen_probs = torch.softmax(gen_logits, dim=-1)
            gen_ids = torch.multinomial(gen_probs.squeeze(0), num_samples=1).squeeze(-1).unsqueeze(0)
            fake_labels = torch.zeros(1, 1).to(device)
            fake_output = discriminator(gen_ids)
            d_fake_loss = criterion(fake_output, fake_labels)
            
            d_loss = d_real_loss + d_fake_loss
            d_loss.backward()
            d_optimizer.step()
            
            g_optimizer.zero_grad()
            gen_logits = generator(noise)
            gen_probs = torch.softmax(gen_logits, dim=-1)
            gen_ids = torch.multinomial(gen_probs.squeeze(0), num_samples=1).squeeze(-1).unsqueeze(0)
            gen_output = discriminator(gen_ids)
            g_loss = criterion(gen_output, real_labels)
            g_loss.backward()
            g_optimizer.step()
        
        print(f"GAN Epoch {epoch+1}/{num_epochs}, D Loss: {d_loss.item():.4f}, G Loss: {g_loss.item():.4f}")

def main():
    global label_list, id2label, all_labels
    
    try:
        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        print(f"使用设备: {device}")
        
        local_model_path = "/root/npl/data/biobert-v1.1"
        train_path = "/root/npl/data/NCBI_corpus_training.txt"
        val_path = "/root/npl/data/NCBI_corpus_validation.txt"
        test_path = "/root/npl/data/NCBI_corpus_testing.txt"
        
        tokenizer = AutoTokenizer.from_pretrained(local_model_path)
        
        print(f"\n读取训练数据集: {train_path}")
        with open(train_path, "r", encoding="utf-8", errors="ignore") as f:
            train_content = f.read()
        
        print(f"\n读取验证数据集: {val_path}")
        with open(val_path, "r", encoding="utf-8", errors="ignore") as f:
            val_content = f.read()
        
        print(f"\n读取测试数据集: {test_path}")
        with open(test_path, "r", encoding="utf-8", errors="ignore") as f:
            test_content = f.read()
        
        print("清洗训练数据集...")
        train_docs, train_labels = clean_dataset(train_content, "训练集")
        
        print("清洗验证数据集...")
        val_docs, val_labels = clean_dataset(val_content, "验证集")
        
        print("清洗测试数据集...")
        test_docs, test_labels = clean_dataset(test_content, "测试集")
        
        if not train_docs:
            print("错误: 训练集清洗后无有效文档")
            return
        
        train_docs = train_docs[:50]
        val_docs = val_docs[:20]
        test_docs = test_docs[:20]
        print(f"使用前 {len(train_docs)} 个训练文档")
        print(f"使用前 {len(val_docs)} 个验证文档")
        print(f"使用前 {len(test_docs)} 个测试文档")
        
        sample = train_docs[0]
        print(f"\n训练集首样本ID: {sample['id']}")
        print(f"文本片段: {sample['text'][:150]}...")
        print(f"实体数量: {len(sample['entities'])}")
        if sample['entities']:
            entity = sample['entities'][0]
            print(f"示例实体: 文本='{entity['text']}', 类型='{entity['label']}', 位置={entity['start']}-{entity['end']}")
        else:
            print("警告: 首样本没有实体")
        
        all_labels = train_labels | val_labels | test_labels
        print(f"\n所有标签类型: {all_labels}")
        label_list = ["O"] + [f"{prefix}-{label}" for prefix in ["B", "I"] for label in sorted(all_labels)]
        label_map = {label: i for i, label in enumerate(label_list)}
        id2label = {i: label for i, label in enumerate(label_list)}
        print(f"创建标签映射: 共{len(label_map)}个标签")
        
        vocab_size = tokenizer.vocab_size
        hidden_dim = 256
        max_length = 128
        generator = Generator(vocab_size, hidden_dim, max_length).to(device)
        discriminator = Discriminator(vocab_size, hidden_dim).to(device)
        
        print("\n训练GAN模型...")
        train_texts = [doc["text"] for doc in train_docs]
        train_gan(generator, discriminator, train_texts, tokenizer, num_epochs=10, device=device)
        
        print("\n生成合成数据...")
        synthetic_texts, synthetic_entities = generate_synthetic_data(
            generator, tokenizer, num_samples=100, max_length=max_length, vocab_size=vocab_size, device=device
        )
        
        synthetic_docs = [
            {"id": f"synthetic_{i}", "text": text, "entities": entities}
            for i, (text, entities) in enumerate(zip(synthetic_texts, synthetic_entities))
        ]
        train_docs.extend(synthetic_docs)
        print(f"合并后训练集大小: {len(train_docs)}")
        
        print("\n准备训练数据集...")
        train_texts = [doc["text"] for doc in train_docs]
        train_entities = [doc["entities"] for doc in train_docs]
        train_dict = {
            "text": train_texts,
            "entities": train_entities
        }
        
        train_dataset = Dataset.from_dict(train_dict)
        tokenized_train = train_dataset.map(
            lambda examples: tokenize_and_align_labels(examples, tokenizer, label_map),
            batched=True
        )
        
        print("\n准备验证数据集...")
        val_texts = [doc["text"] for doc in val_docs]
        val_entities = [doc["entities"] for doc in val_docs]
        val_dict = {
            "text": val_texts,
            "entities": val_entities
        }
        
        val_dataset = Dataset.from_dict(val_dict)
        tokenized_val = val_dataset.map(
            lambda examples: tokenize_and_align_labels(examples, tokenizer, label_map),
            batched=True
        )
        
        print("\n准备测试数据集...")
        test_texts = [doc["text"] for doc in test_docs]
        test_entities = [doc["entities"] for doc in test_docs]
        test_dict = {
            "text": test_texts,
            "entities": test_entities
        }
        
        test_dataset = Dataset.from_dict(test_dict)
        tokenized_test = test_dataset.map(
            lambda examples: tokenize_and_align_labels(examples, tokenizer, label_map),
            batched=True
        )
        
        print(f"训练集大小: {len(tokenized_train)}")
        print(f"验证集大小: {len(tokenized_val)}")
        print(f"测试集大小: {len(tokenized_test)}")
        
        print("\n初始化模型...")
        model = AutoModelForTokenClassification.from_pretrained(
            local_model_path,
            num_labels=len(label_map),
            ignore_mismatched_sizes=True,
            id2label=id2label,
            label2id=label_map
        ).to(device)
        
        training_args = TrainingArguments(
            output_dir="./results",
            num_train_epochs=5,
            per_device_train_batch_size=4,
            per_device_eval_batch_size=4,
            weight_decay=0.01,
            logging_dir='./logs',
            logging_steps=5,
            save_steps=500,
            save_total_limit=1,
            report_to="none",
            learning_rate=5e-5,
            no_cuda=not torch.cuda.is_available(),
            lr_scheduler_type='linear',
            fp16=True,
            eval_steps=500
        )
        
        print("\n初始化训练器...")
        trainer = Trainer(
            model=model,
            args=training_args,
            train_dataset=tokenized_train,
            eval_dataset=tokenized_val,
            compute_metrics=compute_metrics
        )
        
        print("\n开始模型训练...")
        trainer.train()
        
        model.save_pretrained("./results/final_model")
        tokenizer.save_pretrained("./results/final_model")
        print("\n模型训练成功完成并保存!")
        
        print("\n在验证集上评估模型...")
        val_metrics = trainer.evaluate(tokenized_val)
        with open("./results/val_metrics.json", "w") as f:
            json.dump(val_metrics, f, indent=4)
        
        print("\n验证集评估结果:")
        if "eval_overall_precision" in val_metrics:
            print(f"总体准确率: {val_metrics['eval_overall_accuracy']:.4f}")
            print(f"总体精确率: {val_metrics['eval_overall_precision']:.4f}")
            print(f"总体召回率: {val_metrics['eval_overall_recall']:.4f}")
            print(f"总体F1分数: {val_metrics['eval_overall_f1']:.4f}")
        
        val_predictions = trainer.predict(tokenized_val)
        val_cm, val_roc_auc = compute_confusion_matrix_and_roc(val_predictions, label_list, id2label, "验证集")
        
        print("\n验证集混淆矩阵:")
        print(val_cm)
        print("\n验证集每个类别的AUC:")
        for i, label in enumerate(label_list):
            print(f"{label}: {val_roc_auc[i]:.4f}")
        
        for key in val_metrics:
            if "precision" in key or "recall" in key or "f1" in key:
                print(f"{key}: {val_metrics[key]:.4f}")
        
        print("\n在测试集上评估模型...")
        test_metrics = trainer.evaluate(tokenized_test)
        with open("./results/test_metrics.json", "w") as f:
            json.dump(test_metrics, f, indent=4)
        
        print("\n测试集评估结果:")
        if "eval_overall_precision" in test_metrics:
            print(f"总体准确率: {test_metrics['eval_overall_accuracy']:.4f}")
            print(f"总体精确率: {test_metrics['eval_overall_precision']:.4f}")
            print(f"总体召回率: {test_metrics['eval_overall_recall']:.4f}")
            print(f"总体F1分数: {test_metrics['eval_overall_f1']:.4f}")
        
        test_predictions = trainer.predict(tokenized_test)
        test_cm, test_roc_auc = compute_confusion_matrix_and_roc(test_predictions, label_list, id2label, "测试集")
        
        print("\n测试集混淆矩阵:")
        print(test_cm)
        print("\n测试集每个类别的AUC:")
        for i, label in enumerate(label_list):
            print(f"{label}: {test_roc_auc[i]:.4f}")
        
        for key in test_metrics:
            if "precision" in key or "recall" in key or "f1" in key:
                print(f"{key}: {test_metrics[key]:.4f}")
        
        print("\n验证集混淆矩阵已保存至 ./results/验证集_confusion_matrix.png")
        print("验证集ROC曲线已保存至 ./results/验证集_roc_curve.png")
        print("测试集混淆矩阵已保存至 ./results/测试集_confusion_matrix.png")
        print("测试集ROC曲线已保存至 ./results/测试集_roc_curve.png")
        print("\nNLP处理流程完成!")
    
    except Exception as e:
        print(f"\n发生未预期的错误：{str(e)}")
        traceback.print_exc()

if __name__ == "__main__":
    main()

